{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hollweg Dispersion Solver\n",
    "\n",
    "[hollweg]: ../../api/plasmapy.dispersion.numerical.hollweg_.hollweg.rst\n",
    "[bellan2012]: https://doi.org/10.1029/2012ja017856\n",
    "[hollweg1999]: https://doi.org/10.1029/1998ja900132\n",
    "\n",
    "This notebook details the functionality of the [hollweg()][hollweg] function.  This function computes the wave frequencies for given wavenumbers and plasma parameters based on the solution to the two fluid dispersion relation presented by [Hollweg 1999][hollweg1999], and further summarized by [Bellan 2012][bellan2012].  In his derivation Hollweg assumed a uniform magnetic field, zero D.C electric field, quasi-neutrality, and low-frequency waves ($\\omega \\ll \\omega_{ci}$), which yielded the following expression\n",
    "\n",
    "$$\n",
    "    \\left( \\frac{\\omega^2}{{k_z}^2 {v_A}^2} - 1 \\right)\n",
    "    \\left[\\omega^2\n",
    "        \\left(\\omega^2 - k^2 {v_A}^2 \\right)\n",
    "        - \\beta k^2 {v_A}^2 \\left( \\omega^2 - {k_z}^2 {v_A}^2 \\right)\n",
    "    \\right] \\\\\n",
    "    = \\omega^2 \\left(\\omega^2 - k^2 {v_A}^2 \\right) {k_x}^2 \\left(\n",
    "    \\frac{{c_s}^2}{{\\omega_{ci}}^2} - \\frac{c^2}{{\\omega_{pe}}^2}\n",
    "    \\frac{\\omega^2}{{k_z}^2 {v_A}^2} \\right)\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\\beta = c_{s}^2 / v_{A}^2$$\n",
    "\n",
    "$$k_{x} = k \\sin \\theta$$\n",
    "\n",
    "$$k_{z} = k \\cos \\theta$$\n",
    "\n",
    "$$\\mathbf{B_{o}} = B_{o} \\mathbf{\\hat{z}}$$\n",
    "\n",
    "$\\omega$ is the wave frequency, $k$ is the wavenumber\n",
    ", $v_{A}$ is the Alfvén velocity, $c_{s}$ is the ion\n",
    "sound speed, $\\omega_{ci}$ is the ion gyrofrequency, and\n",
    "$\\omega_{pe}$ is the electron plasma frequency.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "Note\n",
    "\n",
    "[Hollweg 1999][hollweg1999] asserts this expression is valid for arbitrary $c_{s} / v_{A}$ and $k_{z} / k$. Contrarily, [Bellan 2012][bellan2012] states in  Section 1.7 that due to the inconsistent retention of the $\\omega / \\omega_{ci} \\ll 1$ terms the expression can only be valid if both $c_{s} \\ll v_{A}$ and the wave propagation is nearly perpendicular to the magnetic field ($|\\theta - \\pi/2| \\ll 0.1$ radians).\n",
    "</div>\n",
    "\n",
    "The [hollweg()][hollweg] function numerically solves for the roots of this equation, which correspond to the Fast, Alfvén, and Acoustic wave modes.\n",
    "\n",
    "## Contents:\n",
    "\n",
    "1. [Wave propagating at nearly 90 degrees](#Wave-propagating-at-nearly-90-degrees)\n",
    "2. [Hollweg 1999 and Bellan 2012 comparison](#Hollweg-1999-and-Bellan-2012-comparison)\n",
    "3. [Reproduce Figure 2 from Hollweg 1999](#Reproduce-Figure-2-from-Hollweg-1999)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import warnings\n",
    "\n",
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib.ticker import MultipleLocator\n",
    "\n",
    "from plasmapy.dispersion.analytical.two_fluid_ import two_fluid\n",
    "from plasmapy.dispersion.numerical.hollweg_ import hollweg\n",
    "from plasmapy.formulary import (\n",
    "    Alfven_speed,\n",
    "    gyrofrequency,\n",
    "    inertial_length,\n",
    "    ion_sound_speed,\n",
    "    plasma_frequency,\n",
    ")\n",
    "from plasmapy.particles import Particle\n",
    "from plasmapy.utils.exceptions import PhysicsWarning\n",
    "\n",
    "warnings.filterwarnings(\n",
    "    action=\"ignore\",\n",
    "    category=PhysicsWarning,\n",
    "    module=\"plasmapy.dispersion.numerical.hollweg_\",\n",
    ")\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = [10.5, 0.56 * 10.5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wave propagating at nearly 90 degrees\n",
    "\n",
    "Below we define the required parameters to compute the wave frequencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define input parameters\n",
    "inputs = {\n",
    "    \"k\": np.logspace(-7, -2, 300) * u.rad / u.m,\n",
    "    \"theta\": 89 * u.deg,\n",
    "    \"n_i\": 5 * u.cm**-3,\n",
    "    \"B\": 1.10232e-8 * u.T,\n",
    "    \"T_e\": 1.6e6 * u.K,\n",
    "    \"T_i\": 4.0e5 * u.K,\n",
    "    \"ion\": Particle(\"p+\"),\n",
    "}\n",
    "\n",
    "# a few useful plasma parameters\n",
    "params = {\n",
    "    \"n_e\": inputs[\"n_i\"] * abs(inputs[\"ion\"].charge_number),\n",
    "    \"cs\": ion_sound_speed(\n",
    "        inputs[\"T_e\"],\n",
    "        inputs[\"T_i\"],\n",
    "        inputs[\"ion\"],\n",
    "    ),\n",
    "    \"va\": Alfven_speed(\n",
    "        inputs[\"B\"],\n",
    "        inputs[\"n_i\"],\n",
    "        ion=inputs[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs[\"B\"], inputs[\"ion\"]),\n",
    "}\n",
    "params[\"lpe\"] = inertial_length(params[\"n_e\"], \"e-\")\n",
    "params[\"wpe\"] = plasma_frequency(params[\"n_e\"], \"e-\")\n",
    "\n",
    "# compute\n",
    "omegas1 = hollweg(**inputs)\n",
    "omegas2 = two_fluid(**inputs)\n",
    "np.set_printoptions(precision=4, threshold=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Quantity]: https://docs.astropy.org/en/stable/api/astropy.units.Quantity.html#astropy.units.Quantity\n",
    "\n",
    "The computed wave frequencies (rad/s) are returned in a dictionary with the keys representing\n",
    "wave modes and the values (instances of Astropy [Quantity]) being the frequencies. Since our inputs were a 1D arrays of\n",
    "of $k$'s, the computed wave frequencies will be a 1D array of size equal to the size of the $k$ array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(list(omegas1.keys()), omegas1[\"fast_mode\"], omegas1[\"fast_mode\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the results of each wave mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 14  # default font size\n",
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 1.6 * figheight\n",
    "fig = plt.figure(figsize=[figwidth, figheight])\n",
    "\n",
    "# normalize data\n",
    "k_prime = inputs[\"k\"] * params[\"lpe\"]\n",
    "\n",
    "# define colormap\n",
    "cmap = plt.get_cmap(\"viridis\")\n",
    "slicedCM = cmap(np.linspace(0, 0.6, 3))\n",
    "\n",
    "# plot\n",
    "(p1,) = plt.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas1[\"fast_mode\"] / params[\"wpe\"]),\n",
    "    \"--\",\n",
    "    c=slicedCM[0],\n",
    "    ms=1,\n",
    "    label=\"Fast\",\n",
    ")\n",
    "ax = plt.gca()\n",
    "(p2,) = ax.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas1[\"alfven_mode\"] / params[\"wpe\"]),\n",
    "    \"--\",\n",
    "    c=slicedCM[1],\n",
    "    ms=1,\n",
    "    label=\"Alfvén\",\n",
    ")\n",
    "(p3,) = ax.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas1[\"acoustic_mode\"] / params[\"wpe\"]),\n",
    "    \"--\",\n",
    "    c=slicedCM[2],\n",
    "    ms=1,\n",
    "    label=\"Acoustic\",\n",
    ")\n",
    "(p4,) = plt.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas2[\"fast_mode\"] / params[\"wpe\"]),\n",
    "    c=slicedCM[0],\n",
    "    ms=1,\n",
    "    label=\"Fast\",\n",
    ")\n",
    "ax = plt.gca()\n",
    "(p5,) = ax.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas2[\"alfven_mode\"] / params[\"wpe\"]),\n",
    "    c=slicedCM[1],\n",
    "    ms=1,\n",
    "    label=\"Alfvén\",\n",
    ")\n",
    "(p6,) = ax.plot(\n",
    "    k_prime,\n",
    "    np.real(omegas2[\"acoustic_mode\"] / params[\"wpe\"]),\n",
    "    c=slicedCM[2],\n",
    "    ms=1,\n",
    "    label=\"Acoustic\",\n",
    ")\n",
    "\n",
    "# adjust axes\n",
    "ax.set_xlabel(r\"$kc / \\omega_{pe}$\", fontsize=fs)\n",
    "ax.set_ylabel(r\"$Re(\\omega / \\omega_{pe})$\", fontsize=fs)\n",
    "ax.set_yscale(\"log\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.tick_params(\n",
    "    which=\"both\",\n",
    "    direction=\"in\",\n",
    "    width=1,\n",
    "    labelsize=fs,\n",
    "    right=True,\n",
    "    length=5,\n",
    ")\n",
    "\n",
    "# annotate\n",
    "styles = [\"-\", \"--\"]\n",
    "s_labels = [\"Hollweg\", \"Bellan\"]\n",
    "ax2 = ax.twinx()\n",
    "for ss, lab in enumerate(styles):\n",
    "    ax2.plot(np.nan, np.nan, ls=lab, label=s_labels[ss], c=\"black\")\n",
    "ax2.get_yaxis().set_visible(False)\n",
    "ax2.legend(fontsize=14, loc=\"lower right\")\n",
    "\n",
    "text1 = (\n",
    "    rf\"$c_s^2/v_A^2 = {params['cs'] ** 2 / params['va'] ** 2:.1f} \\qquad \"\n",
    "    f\"\\\\theta = {inputs['theta'].value:.0f}\"\n",
    "    \"^{\\\\circ}$\"\n",
    ")\n",
    "text2 = (\n",
    "    \"All 3 wave modes are plotted for Hollweg's relation (solid lines) and Bellan's solut- \\n\"\n",
    "    f\"ion (dashed lines) with $\\\\beta = 2.0$ and $\\\\theta = {inputs['theta'].value:.0f}\"\n",
    "    \"^{\\\\circ}$.\"\n",
    ")\n",
    "plt.figtext(-0.08, -0.18, text2, ha=\"left\", transform=ax.transAxes, fontsize=15.5)\n",
    "ax.text(0.57, 0.95, text1, transform=ax.transAxes, fontsize=18)\n",
    "ax.legend(handles=[p4, p1, p5, p2, p6, p3], fontsize=14, ncol=3, loc=\"upper left\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hollweg 1999 and Bellan 2012 comparison\n",
    "\n",
    "[bellan2012]: https://doi.org/10.1029/2012ja017856\n",
    "\n",
    "Figure 1 of [Bellan 2012][bellan2012] chooses parameters such that\n",
    "$\\beta = 0.4$ and $\\Lambda = 0.4$. Below we define parameters to\n",
    "approximate Bellan's assumptions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define input parameters\n",
    "inputs = {\n",
    "    \"B\": 400e-4 * u.T,\n",
    "    \"ion\": Particle(\"He+\"),\n",
    "    \"n_i\": 6.358e19 * u.m**-3,\n",
    "    \"T_e\": 20 * u.eV,\n",
    "    \"T_i\": 10 * u.eV,\n",
    "    \"theta\": np.linspace(0, 90) * u.deg,\n",
    "    \"k\": (2 * np.pi * u.rad) / (0.56547 * u.m),\n",
    "}\n",
    "\n",
    "# a few useful plasma parameters\n",
    "params = {\n",
    "    \"n_e\": inputs[\"n_i\"] * abs(inputs[\"ion\"].charge_number),\n",
    "    \"cs\": ion_sound_speed(inputs[\"T_e\"], inputs[\"T_i\"], inputs[\"ion\"]),\n",
    "    \"wci\": gyrofrequency(inputs[\"B\"], inputs[\"ion\"]),\n",
    "    \"va\": Alfven_speed(inputs[\"B\"], inputs[\"n_i\"], ion=inputs[\"ion\"]),\n",
    "}\n",
    "params[\"beta\"] = (params[\"cs\"] / params[\"va\"]).value ** 2\n",
    "params[\"wpe\"] = plasma_frequency(params[\"n_e\"], \"e-\")\n",
    "params[\"Lambda\"] = (inputs[\"k\"] * params[\"va\"] / params[\"wci\"]).value ** 2\n",
    "\n",
    "(params[\"beta\"], params[\"Lambda\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute omegas for Bellan and Hollweg\n",
    "bellan_omegas = two_fluid(**inputs)\n",
    "hollweg_omegas = hollweg(**inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate data for Bellan curves\n",
    "bellan_plt_vals = {}\n",
    "for mode, arr in bellan_omegas.items():\n",
    "    norm = (np.absolute(arr) / (inputs[\"k\"] * params[\"va\"])).value ** 2\n",
    "    bellan_plt_vals[mode] = {\n",
    "        \"x\": norm * np.sin(inputs[\"theta\"].to(u.rad).value),\n",
    "        \"y\": norm * np.cos(inputs[\"theta\"].to(u.rad).value),\n",
    "    }\n",
    "\n",
    "# generate data for Hollweg curves\n",
    "hollweg_plt_vals = {}\n",
    "for mode, arr in hollweg_omegas.items():\n",
    "    norm = (np.absolute(arr) / (inputs[\"k\"] * params[\"va\"])).value ** 2\n",
    "    hollweg_plt_vals[mode] = {\n",
    "        \"x\": norm * np.sin(inputs[\"theta\"].to(u.rad).value),\n",
    "        \"y\": norm * np.cos(inputs[\"theta\"].to(u.rad).value),\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot all 3 wave modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx-thumbnail": {
     "tooltip": "Hollweg Numerical Dispersion Solver"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fs = 14  # default font size\n",
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 1.6 * figheight\n",
    "fig = plt.figure(figsize=[figwidth, figheight])\n",
    "\n",
    "# define colormap\n",
    "cmap = plt.get_cmap(\"viridis\")\n",
    "slicedCM = cmap(np.linspace(0, 0.6, 3))\n",
    "\n",
    "# Bellan Fast mode\n",
    "(p1,) = plt.plot(\n",
    "    bellan_plt_vals[\"fast_mode\"][\"x\"],\n",
    "    bellan_plt_vals[\"fast_mode\"][\"y\"],\n",
    "    \"--\",\n",
    "    c=slicedCM[0],\n",
    "    linewidth=2,\n",
    "    label=\"Fast\",\n",
    ")\n",
    "ax = plt.gca()\n",
    "\n",
    "# adjust axes\n",
    "ax.set_xlabel(r\"$(\\omega / k v_{A})^{2} \\, \\sin \\theta$\", fontsize=fs)\n",
    "ax.set_ylabel(r\"$(\\omega / k v_{A})^{2} \\, \\cos \\theta$\", fontsize=fs)\n",
    "ax.set_xlim(0.0, 2.0)\n",
    "ax.set_ylim(0.0, 2.0)\n",
    "for spine in ax.spines.values():\n",
    "    spine.set_linewidth(2)\n",
    "ax.minorticks_on()\n",
    "ax.tick_params(which=\"both\", labelsize=fs, width=2)\n",
    "ax.tick_params(which=\"major\", length=10)\n",
    "ax.tick_params(which=\"minor\", length=5)\n",
    "ax.xaxis.set_major_locator(MultipleLocator(0.5))\n",
    "ax.xaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "ax.yaxis.set_major_locator(MultipleLocator(0.5))\n",
    "ax.yaxis.set_minor_locator(MultipleLocator(0.1))\n",
    "\n",
    "# Bellan Alfven mode\n",
    "(p2,) = plt.plot(\n",
    "    bellan_plt_vals[\"alfven_mode\"][\"x\"],\n",
    "    bellan_plt_vals[\"alfven_mode\"][\"y\"],\n",
    "    \"--\",\n",
    "    c=slicedCM[1],\n",
    "    linewidth=2,\n",
    "    label=\"Alfvén\",\n",
    ")\n",
    "\n",
    "# Bellan Acoustic mode\n",
    "(p3,) = plt.plot(\n",
    "    bellan_plt_vals[\"acoustic_mode\"][\"x\"],\n",
    "    bellan_plt_vals[\"acoustic_mode\"][\"y\"],\n",
    "    \"--\",\n",
    "    c=slicedCM[2],\n",
    "    linewidth=2,\n",
    "    label=\"Acoustic\",\n",
    ")\n",
    "\n",
    "# Hollweg Fast mode\n",
    "(p4,) = plt.plot(\n",
    "    hollweg_plt_vals[\"fast_mode\"][\"x\"],\n",
    "    hollweg_plt_vals[\"fast_mode\"][\"y\"],\n",
    "    c=slicedCM[0],\n",
    "    linewidth=2,\n",
    "    label=\"Fast\",\n",
    ")\n",
    "\n",
    "# Hollweg Alfven mode\n",
    "(p5,) = plt.plot(\n",
    "    hollweg_plt_vals[\"alfven_mode\"][\"x\"],\n",
    "    hollweg_plt_vals[\"alfven_mode\"][\"y\"],\n",
    "    c=slicedCM[1],\n",
    "    linewidth=2,\n",
    "    label=\"Alfvén\",\n",
    ")\n",
    "\n",
    "# Hollweg Acoustic mode\n",
    "(p6,) = plt.plot(\n",
    "    hollweg_plt_vals[\"acoustic_mode\"][\"x\"],\n",
    "    hollweg_plt_vals[\"acoustic_mode\"][\"y\"],\n",
    "    c=slicedCM[2],\n",
    "    linewidth=2,\n",
    "    label=\"Acoustic\",\n",
    ")\n",
    "\n",
    "# annotations\n",
    "r = np.linspace(0, 2, 200)\n",
    "X = (r**2) * np.cos(0.1)\n",
    "Y = (r**2) * np.sin(0.1)\n",
    "plt.plot(X, Y, color=\"0.3\")\n",
    "ax.fill_between(X, 0, Y, hatch=\"\\\\\\\\\", color=\"0.7\", alpha=0.5)\n",
    "\n",
    "# style legend\n",
    "styles = [\"-\", \"--\"]\n",
    "s_labels = [\"Hollweg\", \"Bellan\"]\n",
    "ax2 = ax.twinx()\n",
    "for ss, lab in enumerate(styles):\n",
    "    ax2.plot(np.nan, np.nan, ls=lab, label=s_labels[ss], c=\"black\")\n",
    "ax2.get_yaxis().set_visible(False)\n",
    "ax2.legend(fontsize=17, loc=\"center right\")\n",
    "\n",
    "ax.legend(handles=[p4, p1, p5, p2, p6, p3], fontsize=16, ncol=3, loc=\"upper right\")\n",
    "plt.figtext(\n",
    "    1.42,\n",
    "    0.19,\n",
    "    \"$|\\\\theta - \\\\pi / 2| > 0.1 \\\\uparrow$\",\n",
    "    rotation=5.5,\n",
    "    fontsize=20,\n",
    "    transform=ax.transData,\n",
    ")\n",
    "\n",
    "# plot caption\n",
    "txt = (\n",
    "    \"Fig 1. of Bellan 2012 shown with Hollweg's relation (solid lines) and Bellan's anal-\\n\"\n",
    "    \"ytic solution (dashed lines). All 3 wave modes are plotted for $\\\\beta= 0.4$ and $\\\\Lambda= 0.4$.\\n\"\n",
    "    \"The shaded region shows where $\\\\theta \\\\approx \\\\pi / 2$.\\n\"\n",
    ")\n",
    "\n",
    "plt.figtext(-0.1, -0.29, txt, ha=\"left\", transform=ax.transAxes, fontsize=16);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reproduce Figure 2 from Hollweg 1999\n",
    "\n",
    "[hollweg1999]: https://doi.org/10.1029/1998ja900132\n",
    "\n",
    "Figure 2 of [Hollweg 1999][hollweg1999] plots the Alfvén mode\n",
    "and chooses parameters such that $\\beta = 1/20, 1/2, 2, 1/2000$.\n",
    "Below we define parameters to approximate these values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define input parameters\n",
    "# beta = 1/20\n",
    "inputs0 = {\n",
    "    \"k\": np.logspace(-7, -2, 400) * u.rad / u.m,\n",
    "    \"theta\": 90 * u.deg,\n",
    "    \"n_i\": 5 * u.cm**-3,\n",
    "    \"B\": 6.971e-8 * u.T,\n",
    "    \"T_e\": 1.6e6 * u.K,\n",
    "    \"T_i\": 4.0e5 * u.K,\n",
    "    \"ion\": Particle(\"p+\"),\n",
    "}\n",
    "# beta = 1/2\n",
    "inputs1 = {\n",
    "    **inputs0,\n",
    "    \"B\": 2.205e-8 * u.T,\n",
    "}\n",
    "# beta = 2\n",
    "inputs2 = {\n",
    "    **inputs0,\n",
    "    \"B\": 1.10232e-8 * u.T,\n",
    "}\n",
    "# beta = 1/2000\n",
    "inputs3 = {\n",
    "    **inputs0,\n",
    "    \"B\": 6.97178e-7 * u.T,\n",
    "}\n",
    "\n",
    "# a few useful plasma parameters\n",
    "\n",
    "# parameters corresponding to inputs0\n",
    "params0 = {\n",
    "    \"n_e\": inputs0[\"n_i\"] * abs(inputs0[\"ion\"].charge_number),\n",
    "    \"cs\": ion_sound_speed(\n",
    "        inputs0[\"T_e\"],\n",
    "        inputs0[\"T_i\"],\n",
    "        inputs0[\"ion\"],\n",
    "    ),\n",
    "    \"va\": Alfven_speed(\n",
    "        inputs0[\"B\"],\n",
    "        inputs0[\"n_i\"],\n",
    "        ion=inputs0[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs0[\"B\"], inputs0[\"ion\"]),\n",
    "}\n",
    "params0[\"lpe\"] = inertial_length(params0[\"n_e\"], \"e-\")\n",
    "params0[\"wpe\"] = plasma_frequency(params0[\"n_e\"], \"e-\")\n",
    "params0[\"L\"] = params0[\"cs\"] / abs(params0[\"wci\"])\n",
    "\n",
    "# parameters corresponding to inputs1\n",
    "params1 = {\n",
    "    \"n_e\": inputs1[\"n_i\"] * abs(inputs1[\"ion\"].charge_number),\n",
    "    \"cs\": ion_sound_speed(\n",
    "        inputs1[\"T_e\"],\n",
    "        inputs1[\"T_i\"],\n",
    "        inputs1[\"ion\"],\n",
    "    ),\n",
    "    \"va\": Alfven_speed(\n",
    "        inputs1[\"B\"],\n",
    "        inputs1[\"n_i\"],\n",
    "        ion=inputs1[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs1[\"B\"], inputs1[\"ion\"]),\n",
    "}\n",
    "params1[\"lpe\"] = inertial_length(params1[\"n_e\"], \"e-\")\n",
    "params1[\"wpe\"] = plasma_frequency(params1[\"n_e\"], \"e-\")\n",
    "params1[\"L\"] = params1[\"cs\"] / abs(params1[\"wci\"])\n",
    "\n",
    "# parameters corresponding to inputs2\n",
    "params2 = {\n",
    "    \"n_e\": inputs2[\"n_i\"] * abs(inputs2[\"ion\"].charge_number),\n",
    "    \"cs\": ion_sound_speed(\n",
    "        inputs2[\"T_e\"],\n",
    "        inputs2[\"T_i\"],\n",
    "        inputs2[\"ion\"],\n",
    "    ),\n",
    "    \"va\": Alfven_speed(\n",
    "        inputs2[\"B\"],\n",
    "        inputs2[\"n_i\"],\n",
    "        ion=inputs2[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs2[\"B\"], inputs2[\"ion\"]),\n",
    "}\n",
    "params2[\"lpe\"] = inertial_length(params2[\"n_e\"], \"e-\")\n",
    "params2[\"wpe\"] = plasma_frequency(params2[\"n_e\"], \"e-\")\n",
    "params2[\"L\"] = params2[\"cs\"] / abs(params2[\"wci\"])\n",
    "\n",
    "# parameters corresponding to inputs3\n",
    "params3 = {\n",
    "    \"n_e\": inputs3[\"n_i\"] * abs(inputs3[\"ion\"].charge_number),\n",
    "    \"cs\": ion_sound_speed(\n",
    "        inputs3[\"T_e\"],\n",
    "        inputs3[\"T_i\"],\n",
    "        inputs3[\"ion\"],\n",
    "    ),\n",
    "    \"va\": Alfven_speed(\n",
    "        inputs3[\"B\"],\n",
    "        inputs3[\"n_i\"],\n",
    "        ion=inputs3[\"ion\"],\n",
    "    ),\n",
    "    \"wci\": gyrofrequency(inputs3[\"B\"], inputs3[\"ion\"]),\n",
    "}\n",
    "params3[\"lpe\"] = inertial_length(params3[\"n_e\"], \"e-\")\n",
    "params3[\"wpe\"] = plasma_frequency(params3[\"n_e\"], \"e-\")\n",
    "params3[\"L\"] = params3[\"cs\"] / abs(params3[\"wci\"])\n",
    "\n",
    "# confirm beta values\n",
    "beta_vals = [\n",
    "    (params0[\"cs\"] / params0[\"va\"]).value ** 2,\n",
    "    (params1[\"cs\"] / params1[\"va\"]).value ** 2,\n",
    "    (params2[\"cs\"] / params2[\"va\"]).value ** 2,\n",
    "    (params3[\"cs\"] / params3[\"va\"]).value ** 2,\n",
    "]\n",
    "print(\n",
    "    f\"1/{1 / beta_vals[0]:.4f}, \"\n",
    "    f\"1/{1 / beta_vals[1]:.4f}, \"\n",
    "    f\"{beta_vals[2]:.4f}, \"\n",
    "    f\"1/{1 / beta_vals[3]:.4f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 2 of Hollweg 1999 plots over some values that lie outside the valid regime which results in the `PhysicsWarning`'s below being raised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute omegas\n",
    "\n",
    "# show warnings once\n",
    "with warnings.catch_warnings():\n",
    "    warnings.filterwarnings(\"once\")\n",
    "    omegas0 = hollweg(**inputs0)\n",
    "\n",
    "omegas1 = hollweg(**inputs1)\n",
    "omegas2 = hollweg(**inputs2)\n",
    "omegas3 = hollweg(**inputs3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define important quantities for plotting\n",
    "theta = inputs0[\"theta\"].to(u.rad).value\n",
    "\n",
    "kz = np.cos(theta) * inputs0[\"k\"]\n",
    "ky = np.sin(theta) * inputs0[\"k\"]\n",
    "\n",
    "# normalize data\n",
    "k_prime = [\n",
    "    params0[\"L\"] * ky,\n",
    "    params1[\"L\"] * ky,\n",
    "    params2[\"L\"] * ky,\n",
    "    params3[\"L\"] * ky,\n",
    "]\n",
    "big_omega = [\n",
    "    abs(omegas0[\"alfven_mode\"] / (params0[\"va\"] * kz)),\n",
    "    abs(omegas1[\"alfven_mode\"] / (params1[\"va\"] * kz)),\n",
    "    abs(omegas2[\"alfven_mode\"] / (params2[\"va\"] * kz)),\n",
    "    abs(omegas3[\"alfven_mode\"] / (params3[\"va\"] * kz)),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = 14  # default font size\n",
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figheight = 1.6 * figheight\n",
    "fig = plt.figure(figsize=[figwidth, figheight])\n",
    "\n",
    "# define colormap\n",
    "cmap = plt.get_cmap(\"viridis\")\n",
    "slicedCM = cmap(np.linspace(0, 0.6, 4))\n",
    "\n",
    "# plot\n",
    "plt.plot(\n",
    "    k_prime[0],\n",
    "    big_omega[0],\n",
    "    c=slicedCM[0],\n",
    "    ms=1,\n",
    ")\n",
    "ax = plt.gca()\n",
    "ax.plot(\n",
    "    k_prime[1],\n",
    "    big_omega[1],\n",
    "    c=slicedCM[1],\n",
    "    ms=1,\n",
    ")\n",
    "ax.plot(\n",
    "    k_prime[2],\n",
    "    big_omega[2],\n",
    "    c=slicedCM[2],\n",
    "    ms=1,\n",
    ")\n",
    "ax.plot(\n",
    "    k_prime[3],\n",
    "    big_omega[3],\n",
    "    c=slicedCM[3],\n",
    "    ms=1,\n",
    ")\n",
    "\n",
    "# adjust axes\n",
    "ax.set_xlabel(r\"$k_{y}L$\", fontsize=fs)\n",
    "ax.set_ylabel(r\"$|\\Omega|$\", fontsize=fs)\n",
    "ax.set_yscale(\"linear\")\n",
    "ax.set_xscale(\"linear\")\n",
    "ax.set_xlim(0, 1.5)\n",
    "ax.set_ylim(0.96, 1.8)\n",
    "ax.tick_params(\n",
    "    which=\"both\",\n",
    "    direction=\"in\",\n",
    "    width=1,\n",
    "    labelsize=fs,\n",
    "    right=True,\n",
    "    length=5,\n",
    ")\n",
    "\n",
    "# add labels for beta\n",
    "plt.text(1.51, 1.75, \"1/20$=\\\\beta$\", c=slicedCM[0], fontsize=fs)\n",
    "plt.text(1.51, 1.63, \"1/2\", c=slicedCM[1], fontsize=fs)\n",
    "plt.text(1.51, 1.44, \"2\", c=slicedCM[2], fontsize=fs)\n",
    "plt.text(1.51, 0.97, \"1/2000\", c=slicedCM[3], fontsize=fs)\n",
    "\n",
    "# plot caption\n",
    "txt = (\n",
    "    \"Fig. 2 Hollweg 1999 reproduction. The Alfvén mode is plotted for an electron\\n\"\n",
    "    \"-proton plasma with $\\\\beta=$ 1/20, 1/2, 2, 1/2000, $|\\\\Omega|=|\\\\omega / k_{z} v_{A}|$, and $L=c_{s}$ $/ |\\\\omega_{ci}|$.\\n\"\n",
    ")\n",
    "\n",
    "plt.figtext(-0.08, -0.21, txt, ha=\"left\", transform=ax.transAxes, fontsize=16);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "Note \n",
    "\n",
    "Hollweg takes $k_{\\perp}=k_{y}$ for k propagating in the $y$-$z$\n",
    "plane as in the plot above. Contrarily, Bellan takes $k_{\\perp}=k_{x}$\n",
    "for k propagating in the $x$-$z$ plane.\n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "vscode": {
   "interpreter": {
    "hash": "d826483f4fc362efc2f1a43450d731bef6133353e8040b01d718c2f49b5ddeb4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
